import logging
import motor.motor_asyncio
from pymongo import MongoClient, ASCENDING, DESCENDING, TEXT
from pymongo.errors import DuplicateKeyError
from typing import Dict, List, Any, Optional, Union
from datetime import datetime, timedelta
import asyncio
from bson.objectid import ObjectId

class DatabaseManager:
    """
    数据库管理器，负责数据的存储和检索
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化数据库管理器
        
        Args:
            config: 数据库配置
        """
        self.logger = logging.getLogger("database.manager")
        self.config = config.get('database', {})
        
        self.host = self.config.get('host', 'localhost')
        self.port = self.config.get('port', 27017)
        self.db_name = self.config.get('name', 'ai_news_db')
        
        # 初始化MongoDB客户端
        self.client = None
        self.db = None
        self.async_client = None
        self.async_db = None
        
        # 集合配置
        self.collections_config = self.config.get('collections', [])
        
        # 连接数据库
        self._connect()
        
        # 初始化集合和索引
        self._init_collections()
    
    def _connect(self) -> None:
        """
        连接数据库
        """
        try:
            # 同步客户端
            self.client = MongoClient(self.host, self.port)
            self.db = self.client[self.db_name]
            
            # 异步客户端
            connection_string = f"mongodb://{self.host}:{self.port}"
            self.async_client = motor.motor_asyncio.AsyncIOMotorClient(connection_string)
            self.async_db = self.async_client[self.db_name]
            
            self.logger.info(f"Connected to MongoDB at {self.host}:{self.port}")
        except Exception as e:
            self.logger.error(f"Failed to connect to MongoDB: {str(e)}")
            raise
    
    def _init_collections(self) -> None:
        """
        初始化集合和索引
        """
        for collection_config in self.collections_config:
            collection_name = collection_config.get('name')
            if not collection_name:
                continue
                
            try:
                # 创建TTL索引（如果配置了）
                ttl = collection_config.get('ttl')
                if ttl and ttl != 'never':
                    # 解析TTL字符串
                    if ttl.endswith('days'):
                        ttl_seconds = int(ttl.replace('days', '').strip()) * 86400
                    elif ttl.endswith('hours'):
                        ttl_seconds = int(ttl.replace('hours', '').strip()) * 3600
                    else:
                        ttl_seconds = 90 * 86400  # 默认90天
                        
                    self.db[collection_name].create_index(
                        "created_at", 
                        expireAfterSeconds=ttl_seconds
                    )
                
                # 创建其他索引
                if collection_name == 'news':
                    # 新闻集合
                    self.db[collection_name].create_index([("title", TEXT)])
                    self.db[collection_name].create_index("source")
                    self.db[collection_name].create_index("published_date")
                    self.db[collection_name].create_index("topic_categories")
                    self.db[collection_name].create_index("type")
                    self.db[collection_name].create_index("popularity_score")
                
                elif collection_name == 'papers':
                    # 论文集合
                    self.db[collection_name].create_index([("title", TEXT), ("abstract", TEXT)])
                    self.db[collection_name].create_index("authors")
                    self.db[collection_name].create_index("published_date")
                    self.db[collection_name].create_index("citation_count")
                    self.db[collection_name].create_index("categories")
                
                elif collection_name == 'stats':
                    # 统计数据集合
                    self.db[collection_name].create_index("date")
                    self.db[collection_name].create_index("type")
                
                elif collection_name == 'users':
                    # 用户集合
                    self.db[collection_name].create_index("username", unique=True)
                    self.db[collection_name].create_index("email", unique=True)
                
                self.logger.info(f"Initialized collection {collection_name}")
                
            except Exception as e:
                self.logger.error(f"Failed to initialize collection {collection_name}: {str(e)}")
    
    def close(self) -> None:
        """
        关闭数据库连接
        """
        if self.client:
            self.client.close()
        if self.async_client:
            self.async_client.close()
        self.logger.info("Closed database connections")
    
    def save_items(self, collection_name: str, items: List[Dict[str, Any]]) -> int:
        """
        保存多个数据项到指定集合
        
        Args:
            collection_name: 集合名称
            items: 数据项列表
            
        Returns:
            保存成功的数据项数量
        """
        if not items:
            return 0
            
        collection = self.db[collection_name]
        saved_count = 0
        
        for item in items:
            try:
                # 添加创建时间
                if 'created_at' not in item:
                    item['created_at'] = datetime.now()
                
                # 使用更新插入，避免重复
                id_fields = ['id', 'url']
                query = None
                
                for field in id_fields:
                    if field in item and item[field]:
                        query = {field: item[field]}
                        break
                
                if not query:
                    # 如果没有可用的ID字段，使用标题
                    query = {'title': item.get('title')}
                
                result = collection.update_one(
                    query,
                    {'$set': item},
                    upsert=True
                )
                
                if result.upserted_id or result.modified_count > 0:
                    saved_count += 1
                    
            except Exception as e:
                self.logger.error(f"Failed to save item to {collection_name}: {str(e)}")
        
        self.logger.info(f"Saved {saved_count} items to {collection_name}")
        return saved_count
    
    def get_items(self, collection_name: str, query: Dict[str, Any] = None, 
                  sort: List[tuple] = None, limit: int = 100, skip: int = 0) -> List[Dict[str, Any]]:
        """
        从指定集合获取数据
        
        Args:
            collection_name: 集合名称
            query: 查询条件
            sort: 排序规则
            limit: 返回数量限制
            skip: 跳过数量
            
        Returns:
            数据项列表
        """
        collection = self.db[collection_name]
        query = query or {}
        
        try:
            cursor = collection.find(query).skip(skip).limit(limit)
            
            if sort:
                cursor = cursor.sort(sort)
                
            # 将ObjectId转换为字符串
            def convert_objectid(item):
                if '_id' in item and isinstance(item['_id'], ObjectId):
                    item['_id'] = str(item['_id'])
                return item
                
            result = [convert_objectid(item) for item in cursor]
            return result
            
        except Exception as e:
            self.logger.error(f"Failed to get items from {collection_name}: {str(e)}")
            return []
    
    async def save_items_async(self, collection_name: str, items: List[Dict[str, Any]]) -> int:
        """
        异步保存多个数据项到指定集合
        
        Args:
            collection_name: 集合名称
            items: 数据项列表
            
        Returns:
            保存成功的数据项数量
        """
        if not items:
            return 0
            
        collection = self.async_db[collection_name]
        saved_count = 0
        
        for item in items:
            try:
                # 添加创建时间
                if 'created_at' not in item:
                    item['created_at'] = datetime.now()
                
                # 使用更新插入，避免重复
                id_fields = ['id', 'url']
                query = None
                
                for field in id_fields:
                    if field in item and item[field]:
                        query = {field: item[field]}
                        break
                
                if not query:
                    # 如果没有可用的ID字段，使用标题
                    query = {'title': item.get('title')}
                
                result = await collection.update_one(
                    query,
                    {'$set': item},
                    upsert=True
                )
                
                if result.upserted_id or result.modified_count > 0:
                    saved_count += 1
                    
            except Exception as e:
                self.logger.error(f"Failed to save item to {collection_name} async: {str(e)}")
        
        self.logger.info(f"Async saved {saved_count} items to {collection_name}")
        return saved_count
    
    async def get_items_async(self, collection_name: str, query: Dict[str, Any] = None,
                             sort: List[tuple] = None, limit: int = 100, skip: int = 0) -> List[Dict[str, Any]]:
        """
        异步从指定集合获取数据
        
        Args:
            collection_name: 集合名称
            query: 查询条件
            sort: 排序规则
            limit: 返回数量限制
            skip: 跳过数量
            
        Returns:
            数据项列表
        """
        collection = self.async_db[collection_name]
        query = query or {}
        
        try:
            cursor = collection.find(query).skip(skip).limit(limit)
            
            if sort:
                cursor = cursor.sort(sort)
                
            # 将ObjectId转换为字符串
            def convert_objectid(item):
                if '_id' in item and isinstance(item['_id'], ObjectId):
                    item['_id'] = str(item['_id'])
                return item
                
            result = []
            async for item in cursor:
                result.append(convert_objectid(item))
                
            return result
            
        except Exception as e:
            self.logger.error(f"Failed to get items from {collection_name} async: {str(e)}")
            return []
    
    def get_stats(self, start_date: datetime = None, end_date: datetime = None, 
                 group_by: str = 'day', types: List[str] = None) -> List[Dict[str, Any]]:
        """
        获取统计数据
        
        Args:
            start_date: 开始日期
            end_date: 结束日期
            group_by: 分组方式 (day, week, month)
            types: 数据类型列表
            
        Returns:
            统计数据列表
        """
        if not start_date:
            start_date = datetime.now() - timedelta(days=30)
        if not end_date:
            end_date = datetime.now()
            
        pipeline = [
            {'$match': {
                'created_at': {'$gte': start_date, '$lte': end_date}
            }}
        ]
        
        if types:
            pipeline[0]['$match']['type'] = {'$in': types}
        
        # 根据分组方式设置时间格式
        date_format = '%Y-%m-%d'
        if group_by == 'week':
            date_format = '%Y-%U'  # 年份-周数
        elif group_by == 'month':
            date_format = '%Y-%m'  # 年份-月份
            
        # 添加分组聚合
        pipeline.extend([
            {'$project': {
                'date': {'$dateToString': {'format': date_format, 'date': '$created_at'}},
                'type': 1,
                'topic_categories': 1
            }},
            {'$group': {
                '_id': {
                    'date': '$date',
                    'type': '$type'
                },
                'count': {'$sum': 1},
                'categories': {'$push': '$topic_categories'}
            }},
            {'$sort': {'_id.date': 1}}
        ])
        
        try:
            # 执行聚合查询
            result = list(self.db['news'].aggregate(pipeline))
            
            # 格式化结果
            formatted_results = []
            for item in result:
                # 展平类别
                all_categories = []
                for categories in item.get('categories', []):
                    if isinstance(categories, list):
                        all_categories.extend(categories)
                    
                # 统计类别频率
                category_counts = {}
                for category in all_categories:
                    if category in category_counts:
                        category_counts[category] += 1
                    else:
                        category_counts[category] = 1
                        
                formatted_item = {
                    'date': item['_id']['date'],
                    'type': item['_id']['type'],
                    'count': item['count'],
                    'category_stats': category_counts
                }
                formatted_results.append(formatted_item)
                
            return formatted_results
            
        except Exception as e:
            self.logger.error(f"Failed to get stats: {str(e)}")
            return []
    
    def save_stats(self, stats: List[Dict[str, Any]]) -> int:
        """
        保存统计数据
        
        Args:
            stats: 统计数据列表
            
        Returns:
            保存成功的数据项数量
        """
        return self.save_items('stats', stats)
    
    def search_items(self, query_text: str, collection_name: str = 'news', 
                    filters: Dict[str, Any] = None, limit: int = 20) -> List[Dict[str, Any]]:
        """
        搜索数据
        
        Args:
            query_text: 搜索文本
            collection_name: 集合名称
            filters: 过滤条件
            limit: 返回数量限制
            
        Returns:
            搜索结果列表
        """
        collection = self.db[collection_name]
        
        try:
            # 构造文本搜索查询
            text_query = {'$text': {'$search': query_text}}
            
            # 合并过滤条件
            if filters:
                query = {'$and': [text_query, filters]}
            else:
                query = text_query
                
            # 添加文本相关性评分
            projection = {'score': {'$meta': 'textScore'}}
            
            # 执行查询
            cursor = collection.find(query, projection).sort([('score', {'$meta': 'textScore'})]).limit(limit)
            
            # 将ObjectId转换为字符串
            def convert_objectid(item):
                if '_id' in item and isinstance(item['_id'], ObjectId):
                    item['_id'] = str(item['_id'])
                return item
                
            result = [convert_objectid(item) for item in cursor]
            return result
            
        except Exception as e:
            self.logger.error(f"Failed to search items in {collection_name}: {str(e)}")
            return [] 